#!/bin/bash

# ============================================
# LoRA 两阶段训练 - 阶段2：只训练下游任务层
# ============================================
# 此阶段加载阶段1的 LoRA 检查点，冻结 LoRA，只训练下游任务层
# 
# 使用方法：
# 1. 先运行 stage1 脚本训练 LoRA
# 2. 将 stage1 的检查点路径设置为 STAGE1_CHECKPOINT 变量
# 3. 运行此脚本

DATASET_NAME="PPI"
DATASET_TYPE="protein"
TASK_TYPE="binary_class"
TASK_LEVEL_TYPE="seq_level"
MODEL_TYPE="lucappi"
CONFIG_NAME="ppi_config.json"
INPUT_TYPE="matrix"
INPUT_MODE="pair"
LABEL_TYPE="PPI"
FUSION_TYPE="concat"
embedding_input_size=2560
SEQ_MAX_LENGTH=2048
matrix_max_length=2048
TRUNC_TYPE="right"
hidden_size=1024
num_attention_heads=8
num_hidden_layers=4
dropout_prob=0.1
SEQ_POOLING_TYPE="value_attention"
MATRIX_POOLING_TYPE="value_attention"
VOCAB_NAME="gene_prot"
BEST_METRIC_TYPE="acc"
classifier_size=128
loss_type="bce"

llm_type="lucaone"
llm_version="lucaone"
llm_step=36000000

# ========== LoRA 参数（必须与阶段1一致）==========
LORA_RANK=8
LORA_ALPHA=16.0
LORA_DROPOUT=0.1
LORA_TARGET_MODULES="self_attn.q_proj,self_attn.k_proj,self_attn.v_proj,self_attn.out_proj,fc1,fc2"

# ========== 阶段2 训练配置 ==========
# 阶段2：只训练下游任务层，使用较大的学习率
batch_size=16
matrix_max_length=1024
nproc_per_node=8
learning_rate=1e-3                      # 下游任务层使用较大学习率
gradient_accumulation_steps=8
buffer_size=2048
worker_num=0
num_train_epochs=20                     # 阶段2 训练轮数（可根据需要调整）
time_str=$(date "+%Y%m%d%H%M%S")
STAGE="stage2_downstream_only"           # 阶段标识

# ========== 重要：设置阶段1的检查点路径 ==========
# 请将下面的路径替换为阶段1实际保存的检查点路径
# 例如：STAGE1_CHECKPOINT="../models/PPI/protein/binary_class/lucappi_lora_stage1_lora_only/matrix/20240101120000"
STAGE1_CHECKPOINT=""                     # 请填写阶段1的检查点路径

# 检查是否设置了检查点路径
if [ -z "$STAGE1_CHECKPOINT" ]; then
    echo "错误：请先设置 STAGE1_CHECKPOINT 变量为阶段1的检查点路径！"
    echo "例如：STAGE1_CHECKPOINT=\"../models/PPI/protein/binary_class/lucappi_lora_stage1_lora_only/matrix/20240101120000\""
    exit 1
fi

# 检查检查点是否存在
if [ ! -d "$STAGE1_CHECKPOINT" ]; then
    echo "错误：检查点路径不存在: $STAGE1_CHECKPOINT"
    exit 1
fi

cd ../../

echo "============================================"
echo "阶段2：只训练下游任务层"
echo "============================================"
echo "训练配置："
echo "  - 冻结 LoRA: True"
echo "  - 只训练下游任务层: True"
echo "  - 学习率: $learning_rate"
echo "  - 训练轮数: $num_train_epochs"
echo "  - 加载 LoRA 检查点: $STAGE1_CHECKPOINT"
echo "============================================"

# 使用 torchrun 进行多 GPU 分布式训练
torchrun --nproc_per_node=$nproc_per_node run.py \
  --train_data_dir ../dataset/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/train/ \
  --dev_data_dir ../dataset/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/dev/ \
  --test_data_dir ../dataset/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/test/ \
  --dataset_name $DATASET_NAME \
  --dataset_type $DATASET_TYPE \
  --task_type $TASK_TYPE \
  --task_level_type $TASK_LEVEL_TYPE \
  --model_type $MODEL_TYPE \
  --input_type $INPUT_TYPE \
  --input_mode $INPUT_MODE \
  --label_type $LABEL_TYPE \
  --alphabet gene_prot \
  --label_filepath ../dataset/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/label.txt  \
  --output_dir ../models/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/${MODEL_TYPE}_lora_${STAGE}/$INPUT_TYPE/$time_str \
  --log_dir ../logs/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/${MODEL_TYPE}_lora_${STAGE}/$INPUT_TYPE/$time_str \
  --tb_log_dir ../tb-logs/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/${MODEL_TYPE}_lora_${STAGE}/$INPUT_TYPE/$time_str \
  --config_path ../config/$MODEL_TYPE/$CONFIG_NAME \
  --seq_vocab_path  gene_prot \
  --seq_pooling_type $SEQ_POOLING_TYPE \
  --matrix_pooling_type $MATRIX_POOLING_TYPE \
  --fusion_type $FUSION_TYPE \
  --do_train \
  --do_eval \
  --do_predict \
  --do_metrics \
  --evaluate_during_training \
  --per_gpu_train_batch_size=$batch_size \
  --per_gpu_eval_batch_size=$batch_size \
  --gradient_accumulation_steps=$gradient_accumulation_steps \
  --learning_rate=$learning_rate \
  --lr_update_strategy step \
  --lr_decay_rate 0.9 \
  --num_train_epochs=$num_train_epochs \
  --overwrite_output_dir \
  --seed 1221 \
  --sigmoid \
  --loss_type $loss_type \
  --best_metric_type $BEST_METRIC_TYPE \
  --seq_max_length=$SEQ_MAX_LENGTH \
  --embedding_input_size $embedding_input_size \
  --matrix_max_length=$matrix_max_length \
  --trunc_type=$TRUNC_TYPE \
  --no_token_embeddings \
  --no_token_type_embeddings \
  --no_position_embeddings \
  --pos_weight 1.0 \
  --buffer_size $buffer_size \
  --worker_num $worker_num \
  --delete_old \
  --llm_dir .. \
  --llm_type $llm_type \
  --llm_version $llm_version \
  --llm_step $llm_step \
  --ignore_index -100 \
  --hidden_size $hidden_size \
  --num_attention_heads $num_attention_heads \
  --num_hidden_layers $num_hidden_layers \
  --dropout_prob $dropout_prob \
  --classifier_size $classifier_size \
  --vector_dirpath ../../vectors/$DATASET_NAME/$llm_type/$llm_version/$llm_step   \
  --matrix_dirpath ../../matrices/$DATASET_NAME/$llm_type/$llm_version/$llm_step  \
  --seq_fc_size null \
  --matrix_fc_size 128 \
  --vector_fc_size null \
  --emb_activate_func gelu \
  --fc_activate_func gelu \
  --classifier_activate_func gelu \
  --warmup_steps 1000 \
  --beta1 0.9 \
  --beta2 0.98 \
  --weight_decay 0.01 \
  --save_steps -1 \
  --logging_steps 200 \
  --matrix_encoder \
  --fp16 \
  --fp16_opt_level O1 \
  --use_lora \
  --lora_rank $LORA_RANK \
  --lora_alpha $LORA_ALPHA \
  --lora_dropout $LORA_DROPOUT \
  --lora_target_modules "$LORA_TARGET_MODULES" \
  --lora_checkpoint $STAGE1_CHECKPOINT \
  --freeze_lora

echo ""
echo "============================================"
echo "阶段2 训练完成！"
echo "最终模型保存在: ../models/$DATASET_NAME/$DATASET_TYPE/$TASK_TYPE/${MODEL_TYPE}_lora_${STAGE}/$INPUT_TYPE/$time_str"
echo "============================================"

